from openpyxl import load_workbook
import re
from docx import Document
from docxtpl import DocxTemplate
import logging
from typing import Dict, Any, List, Tuple, Optional
import gc

class TemplateEngine:
    def __init__(self, custom_var_extractor=None):
        """
        初始化 TemplateEngine 实例。

        :param template_path: Word 模板文件的路径 (str)
        :param excel_file_path: Excel 文件的路径 (str)
        :param output_path: 生成的 Word 文档保存路径 (str)
        :param sheet_name: Excel 工作表的名称，默认为 'Sheet1' (str)
        :param custom_var_extractor: 用于提取自定义变量的回调函数或方法 (callable or None)
        """
        self.template_path = None
        self.excel_file_path = None
        self.output_path = None
        self.sheet_name = None
        self.custom_var_extractor = custom_var_extractor
        self.doc=None
        
        # 初始化属性
        self.placeholders = {}
        self.data = {}
    
    def load_template(self, template_path: str) -> Tuple[bool, List[str]]:
        """加载Word模板"""
        try:
            self.cleanup()
            self.template_path = template_path

            variables = self.extract_placeholders_from_docx()
            return True, variables
        except Exception as e:
            self.cleanup()
            return False, [f"加载模板失败: {str(e)}"]
    
    def load_excel(self, excel_file_path: str) -> Tuple[bool, List[str]]:
        """加载Excel"""
        try:
            self.excel_file_path = excel_file_path  # 统一使用 self.excel_file_path
            self._get_data_from_excel()
            return True, [f"加载Excel成功"]
        except Exception as e:
            self.cleanup()
            return False, [f"加载EXCEL失败: {str(e)}"]

    def _extract_placeholders(self, elements, pattern):
        """从给定的文档元素中提取所有匹配正则表达式的占位符，并分类存储"""
        for element in elements:
            matches = pattern.findall(element.text) if hasattr(element, 'text') else []
            for match in matches:
                placeholder_type, var_name = match
                if placeholder_type not in self.placeholders:
                    self.placeholders[placeholder_type] = set()
                self.placeholders[placeholder_type].add(var_name)

    def extract_placeholders_from_docx(self):
        """从Word文档中提取所有Jinja2风格的占位符，并分类存储在 self.placeholders 中"""
        # 定义正则表达式来匹配不同类型的占位符
        placeholder_pattern = re.compile(r'{{\s*(custom_|cell_)?([a-zA-Z0-9_]+)\s*}}')

        # 打开Word文档
        doc = Document(self.template_path)
       
        # 提取段落中的占位符
        self._extract_placeholders(doc.paragraphs, placeholder_pattern)
        
        # 提取表格中的占位符
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    self._extract_placeholders([cell], placeholder_pattern)
        
        
        return self.placeholders

    def _get_data_from_excel(self, sheet_name: Optional[str] = None):
        """根据变量名从Excel文件中读取特定单元格或区域的数据，并存储在 self.data 中"""
        # 定义正则表达式来匹配 `cell_` 格式的变量名
        cell_pattern = re.compile(r'cell_((?P<col_start>[A-Za-z]+)(?P<row_start>\d+))(_(?P<col_end>[A-Za-z]+)(?P<row_end>\d+))?')
        
        try:
            # 使用 openpyxl 加载 Excel 文件
            wb = load_workbook(self.excel_file_path)

            sheet_names = wb.sheetnames
            
            if not sheet_names:
                logging.error("The workbook is empty.")
                raise ValueError("The workbook is empty.")
            
            # 获取默认第一张工作表或指定的工作表
            ws = wb[sheet_names[0]] if sheet_name is None else wb[sheet_name]
                
            for var in self.placeholders.get('cell_', set()):
                match = cell_pattern.match(f"cell_{var}")
                if match:
                    col_start = match.group('col_start')
                    row_start = int(match.group('row_start'))
                    col_end = match.group('col_end')
                    row_end = int(match.group('row_end')) if match.group('row_end') else None
                    
                    try:
                        if col_end and row_end:
                            # 区域范围
                            values = []
                            for row in ws[f"{col_start}{row_start}:{col_end}{row_end}"]:
                                row_values = [cell.value for cell in row]
                                values.append(row_values)
                            self.data[f"cell_{var}"] = values
                        else:
                            # 单个单元格
                            value = ws[f"{col_start}{row_start}"].value
                            self.data[f"cell_{var}"] = value
                    except KeyError as e:
                        logging.error(f"Column '{col_start}' or '{col_end}' not found in the Excel file.")
                        self.data[f"cell_{var}"] = None
                else:
                    logging.warning(f"Variable 'cell_{var}' does not match the expected format.")
                    self.data[f"cell_{var}"] = None
            return self.data
        except Exception as e:
            logging.error(f"Error reading Excel file: {e}", exc_info=True)  # 记录完整的错误栈跟踪信息
            raise ValueError(f"Error reading Excel file: {e}")

    def extract_custom_variables(self):
        """从外部数据源提取自定义变量，并存储在 self.data 中"""
        custom_vars = self.placeholders.get('custom_', set())
        if not custom_vars:
            logging.info("No custom variables found in the template.")
            return
        
        if callable(self.custom_var_extractor):
            try:
                custom_data = self.custom_var_extractor(custom_vars)
                if not isinstance(custom_data, dict):
                    raise ValueError("Custom variable extractor must return a dictionary.")
                self.data.update({f"custom_{k}": v for k, v in custom_data.items()})
            except Exception as e:
                logging.error(f"Error extracting custom variables: {e}")
        else:
            logging.warning("Custom variable extractor is not provided or not callable.")

    def save_document(self, output_path: str) -> Tuple[bool, str]:
        """保存文档"""
        try:
            if not self.template_path:
                return False, "未加载模板"
            
            if not self.placeholders:
                return False, "未找到占位符"
            
            if not self.data:
                return False, "未找到数据"
            
            self.extract_custom_variables()
            
            # 加载Word模板
            self.doc = DocxTemplate(self.template_path)
            
            # 准备要插入的数据 (dict)
            context = {
                f"{prefix}{var}": self.data.get(f"{prefix}{var}") 
                for prefix, vars_set in self.placeholders.items() 
                for var in vars_set
            }
            
            # 渲染文档
            self.doc.render(context)
            
            # 保存生成的文档
            self.doc.save(output_path)  # 使用传入的 output_path 而不是 self.output_path
            self.output_path = output_path  # 更新 self.output_path
            
            logging.info(f"Document saved at: {output_path}")
            return True, "保存成功"
        except Exception as e:
            return False, f"保存失败: {str(e)}"
        finally:
            self.cleanup()
    def cleanup(self) -> None:
        """清理资源但保留路径"""
        try:
            if hasattr(self, 'doc') and self.doc:
                self.doc = None
            
            # 不清理template_path，保留路径信息
            
            # 强制进行垃圾回收
            gc.collect()
        except Exception as e:
            print(f"清理资源时出错: {str(e)}") 